

package com.hazelcast.collection.impl.queue;

import com.hazelcast.collection.QueueStore;
import com.hazelcast.collection.QueueStoreFactory;
import com.hazelcast.config.QueueStoreConfig;
import com.hazelcast.internal.diagnostics.Diagnostics;
import com.hazelcast.internal.diagnostics.StoreLatencyPlugin;
import com.hazelcast.internal.namespace.NamespaceUtil;
import com.hazelcast.internal.nio.ClassLoaderUtil;
import com.hazelcast.internal.serialization.Data;
import com.hazelcast.internal.serialization.SerializationService;
import com.hazelcast.internal.serialization.impl.HeapData;
import com.hazelcast.internal.util.StringUtil;
import com.hazelcast.spi.impl.NodeEngine;
import com.hazelcast.spi.impl.NodeEngineImpl;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.*;

import static com.hazelcast.config.QueueStoreConfig.DEFAULT_BULK_LOAD;
import static com.hazelcast.config.QueueStoreConfig.DEFAULT_MEMORY_LIMIT;
import static com.hazelcast.internal.util.EmptyStatement.ignore;
import static com.hazelcast.internal.util.MapUtil.createHashMap;
import static com.hazelcast.internal.util.Preconditions.checkNotNull;

/**
 * Wrapper for the Queue Store.
 */
public final class QueueStoreWrapper implements QueueStore<Data> {
    private final String name;
    private final @Nullable String namespace;
    private final NodeEngine nodeEngine;
    private int memoryLimit = DEFAULT_MEMORY_LIMIT;
    private int bulkLoad = DEFAULT_BULK_LOAD;
    private boolean enabled;
    private boolean binary;
    private QueueStore store;
    private SerializationService serializationService;

    private QueueStoreWrapper(@Nonnull NodeEngine nodeEngine, @Nonnull String name, @Nullable String namespace) {
        this.name = name;
        this.nodeEngine = nodeEngine;
        this.namespace = namespace;
    }

    /**
     * Factory method that creates a {@link QueueStoreWrapper}
     *
     * @param name                 queue name
     * @param storeConfig          store config of queue
     * @param serializationService serialization service.
     * @return returns a new instance of {@link QueueStoreWrapper}
     */
    public static QueueStoreWrapper create(@Nonnull NodeEngine nodeEngine, @Nonnull String name, @Nullable QueueStoreConfig storeConfig, @Nonnull SerializationService serializationService, @Nullable ClassLoader classLoader, @Nullable String namespace) {
        checkNotNull(name, "name should not be null");
        checkNotNull(serializationService, "serializationService should not be null");

        QueueStoreWrapper storeWrapper = new QueueStoreWrapper(nodeEngine, name, namespace);
        storeWrapper.setSerializationService(serializationService);
        if (storeConfig == null || !storeConfig.isEnabled()) {
            return storeWrapper;
        }
        // create queue store.
        QueueStore queueStore = createQueueStore(name, storeConfig, classLoader);
        if (queueStore != null) {
            boolean isBinary = Boolean.parseBoolean(storeConfig.getProperty(QueueStoreConfig.STORE_BINARY));
            int memoryLimit = parseInt(QueueStoreConfig.STORE_MEMORY_LIMIT, DEFAULT_MEMORY_LIMIT, storeConfig);
            int bulkLoad = parseInt(QueueStoreConfig.STORE_BULK_LOAD, DEFAULT_BULK_LOAD, storeConfig);
            storeWrapper.setEnabled(storeConfig.isEnabled()).setBinary(isBinary).setMemoryLimit(memoryLimit).setBulkLoad(bulkLoad).setStore(queueStore);
        }
        return storeWrapper;
    }

    private static QueueStore createQueueStore(String name, QueueStoreConfig storeConfig, ClassLoader classLoader) {
        // 1. Try to create store from `store impl.` class.
        QueueStore store = getQueueStore(storeConfig, classLoader);
        // 2. Try to create store from `store factory impl.` class.
        if (store == null) {
            store = getQueueStoreFactory(name, storeConfig, classLoader);
        }
        return store;
    }

    private static QueueStore getQueueStore(QueueStoreConfig storeConfig, ClassLoader classLoader) {
        if (storeConfig == null) {
            return null;
        }
        QueueStore store = storeConfig.getStoreImplementation();
        if (store != null) {
            return store;
        }
        try {
            store = ClassLoaderUtil.newInstance(classLoader, storeConfig.getClassName());
        } catch (Exception ignored) {
            ignore(ignored);
        }
        return store;
    }

    private static QueueStore getQueueStoreFactory(String name, QueueStoreConfig storeConfig, ClassLoader classLoader) {
        if (storeConfig == null) {
            return null;
        }
        QueueStoreFactory factory = storeConfig.getFactoryImplementation();
        if (factory == null) {
            try {
                factory = ClassLoaderUtil.newInstance(classLoader, storeConfig.getFactoryClassName());
            } catch (Exception ignored) {
                ignore(ignored);
            }
        }
        return factory == null ? null : factory.newQueueStore(name, storeConfig.getProperties());
    }

    private static int parseInt(String name, int defaultValue, QueueStoreConfig storeConfig) {
        String val = storeConfig.getProperty(name);
        if (StringUtil.isNullOrEmptyAfterTrim(val)) {
            return defaultValue;
        }
        try {
            return Integer.parseInt(val);
        } catch (NumberFormatException e) {
            return defaultValue;
        }
    }

    void instrument(NodeEngine nodeEngine) {
        Diagnostics diagnostics = ((NodeEngineImpl) nodeEngine).getDiagnostics();
        StoreLatencyPlugin storeLatencyPlugin = diagnostics.getPlugin(StoreLatencyPlugin.class);
        if (!enabled || storeLatencyPlugin == null) {
            return;
        }

        this.store = new LatencyTrackingQueueStore(store, storeLatencyPlugin, name);
    }

    @Override
    public void store(Long key, Data value) {
        if (!enabled) {
            return;
        }
        NamespaceUtil.runWithNamespace(nodeEngine, namespace, () -> {
            Object actualValue;
            if (binary) {
                // WARNING: we can't pass original Data to the user
                actualValue = Arrays.copyOf(value.toByteArray(), value.totalSize());
            } else {
                actualValue = serializationService.toObject(value);
            }
            store.store(key, actualValue);
        });
    }

    @Override
    public void storeAll(Map<Long, Data> map) {
        if (!enabled) {
            return;
        }

        NamespaceUtil.runWithNamespace(nodeEngine, namespace, () -> {
            Map<Long, Object> objectMap = createHashMap(map.size());
            if (binary) {
                // WARNING: we can't pass original Data to the user
                // TODO: @mm - is there really an advantage of using binary storeAll?
                // since we need to do array copy for each item.
                for (Map.Entry<Long, Data> entry : map.entrySet()) {
                    Data value = entry.getValue();
                    byte[] copy = Arrays.copyOf(value.toByteArray(), value.totalSize());
                    objectMap.put(entry.getKey(), copy);
                }
            } else {
                for (Map.Entry<Long, Data> entry : map.entrySet()) {
                    objectMap.put(entry.getKey(), serializationService.toObject(entry.getValue()));
                }
            }
            store.storeAll(objectMap);
        });
    }

    @Override
    public void delete(Long key) {
        if (enabled) {
            NamespaceUtil.runWithNamespace(nodeEngine, namespace, () -> store.delete(key));
        }
    }

    @Override
    public void deleteAll(Collection<Long> keys) {
        if (enabled) {
            NamespaceUtil.runWithNamespace(nodeEngine, namespace, () -> store.deleteAll(keys));
        }
    }

    @Override
    public Data load(Long key) {
        if (!enabled) {
            return null;
        }

        return NamespaceUtil.callWithNamespace(nodeEngine, namespace, () -> {
            Object val = store.load(key);
            if (binary) {
                byte[] dataBuffer = (byte[]) val;
                return new HeapData(Arrays.copyOf(dataBuffer, dataBuffer.length));
            }
            return serializationService.toData(val);
        });
    }

    @Override
    public Map<Long, Data> loadAll(Collection<Long> keys) {
        if (!enabled) {
            return null;
        }

        return NamespaceUtil.callWithNamespace(nodeEngine, namespace, () -> {
            Map<Long, ?> map = store.loadAll(keys);
            if (map == null) {
                return Collections.emptyMap();
            }
            Map<Long, Data> dataMap = createHashMap(map.size());
            if (binary) {
                for (Map.Entry<Long, ?> entry : map.entrySet()) {
                    byte[] dataBuffer = (byte[]) entry.getValue();
                    Data data = new HeapData(Arrays.copyOf(dataBuffer, dataBuffer.length));
                    dataMap.put(entry.getKey(), data);
                }
            } else {
                for (Map.Entry<Long, ?> entry : map.entrySet()) {
                    dataMap.put(entry.getKey(), serializationService.toData(entry.getValue()));
                }
            }
            return dataMap;
        });
    }

    @Override
    public Set<Long> loadAllKeys() {
        if (enabled) {
            return NamespaceUtil.callWithNamespace(nodeEngine, namespace, () -> store.loadAllKeys());
        }
        return null;
    }

    public boolean isEnabled() {
        return enabled;
    }

    QueueStoreWrapper setEnabled(boolean enabled) {
        this.enabled = enabled;
        return this;
    }

    public boolean isBinary() {
        return binary;
    }

    QueueStoreWrapper setBinary(boolean binary) {
        this.binary = binary;
        return this;
    }

    public int getMemoryLimit() {
        return memoryLimit;
    }

    QueueStoreWrapper setMemoryLimit(int memoryLimit) {
        this.memoryLimit = memoryLimit;
        return this;
    }

    public int getBulkLoad() {
        return bulkLoad;
    }

    QueueStoreWrapper setBulkLoad(int bulkLoad) {
        if (bulkLoad < 1) {
            bulkLoad = 1;
        }
        this.bulkLoad = bulkLoad;
        return this;
    }

    void setSerializationService(SerializationService serializationService) {
        this.serializationService = serializationService;
    }

    QueueStoreWrapper setStore(QueueStore store) {
        this.store = store;
        return this;
    }
}
